{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pinecone-io/examples/blob/master/learn/generation/chatbots/nemo-guardrails/01-variables-and-flows.ipynb) [![Open nbviewer](https://raw.githubusercontent.com/pinecone-io/examples/master/assets/nbviewer-shield.svg)](https://nbviewer.org/github/pinecone-io/examples/blob/master/learn/generation/chatbots/nemo-guardrails/01-variables-and-flows.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "r1o27sM3we_3"
      },
      "outputs": [],
      "source": [
        "!pip install -qU \\\n",
        "    nemoguardrails==0.4.0 \\\n",
        "    openai==0.27.8"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "pO2LkrOb2i74"
      },
      "source": [
        "# Variables and Flows\n",
        "\n",
        "Colang also allows us to use logical flows (if-else) and insert variables into our rails. There are several ways of inputing variables. When inserting the full conversational memory into our rails (rather than a single `prompt` as we have been doing so far) we can use the `context` message to initialize variables.\n",
        "\n",
        "Let's begin by doing that. First, we'll initialize a simple Rail."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "IX9eHPvKxJHK"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = os.environ.get(\"OPENAI_API_KEY\") or \"YOUR_API_KEY\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "DPjxD_5mxaTK"
      },
      "outputs": [],
      "source": [
        "colang_content = \"\"\"\n",
        "define user greeting\n",
        "    \"Hey there!\"\n",
        "    \"How are you?\"\n",
        "    \"What's up?\"\n",
        "\n",
        "define bot name greeting\n",
        "    \"Hey $name!\"\n",
        "\n",
        "define flow\n",
        "    user greeting\n",
        "    if $name\n",
        "        bot name greeting\n",
        "    else\n",
        "        bot greeting\n",
        "\"\"\"\n",
        "yaml_content = \"\"\"\n",
        "models:\n",
        "- type: main\n",
        "  engine: openai\n",
        "  model: text-davinci-003\n",
        "\"\"\""
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "SQsM7utD3SKl"
      },
      "source": [
        "Initialize the Rail:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "yu_R040G3TdI"
      },
      "outputs": [],
      "source": [
        "from nemoguardrails import LLMRails, RailsConfig\n",
        "\n",
        "# initialize rails config\n",
        "config = RailsConfig.from_content(\n",
        "    colang_content=colang_content,\n",
        "    yaml_content=yaml_content\n",
        ")\n",
        "# create rails\n",
        "rails = LLMRails(config)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "q9BLM2he3Uyj"
      },
      "source": [
        "Now, to begin using the `messages` input we need to set a list of messages which must be structured as a list of dictionary objects each containing `\"role\"` and `\"content\"` keys. Guardrails allows us to use three roles, `\"context\"`, `\"user\"`, and `\"assistant\"`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "bqsPOCZM3p36"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\"role\": \"context\", \"content\": \"\"},\n",
        "    {\"role\": \"user\", \"content\": \"Hey there!\"}\n",
        "    #{\"role\": \"assistant\",\n",
        "    # \"content\": \"Hi! How are you? Is there anything I can help with?\"},\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aVcKmJ685FV9",
        "outputId": "2dc4c64b-1d88-41be-8fe8-4edd2be0aa60"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'role': 'assistant', 'content': 'Hey there! How can I help you?'}"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "await rails.generate_async(messages=messages)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Zp8g2eX-_Gh2"
      },
      "source": [
        "Right now the bot doesn't greet us with our name because we left the `context` empty. Instead, let's try passing in the `$name` parameter."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "ePFrXLBT5V2e"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\"role\": \"context\", \"content\": {\"name\": \"James\"}},\n",
        "    {\"role\": \"assistant\",\n",
        "     \"content\": \"Hi! How are you? Is there anything I can help with?\"},\n",
        "    {\"role\": \"user\", \"content\": \"Hey there!\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZgUXe3dg8nlo",
        "outputId": "c2607df8-3e9a-405b-a79c-1cfc56a5f7ea"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'role': 'assistant', 'content': 'Hey James!'}"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "await rails.generate_async(messages=messages)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "2Ocz_xZw5HGG"
      },
      "source": [
        "Above we're using the `context` to set context variables within our Rails. Sometimes we may want to set them from within a conversation. Let's see how we could grab the users name and pass it into a `$name` variable."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "mEqWmQLB5ba9"
      },
      "outputs": [],
      "source": [
        "colang_content = \"\"\"\n",
        "define user give name\n",
        "    \"My name is James\"\n",
        "    \"I'm Julio\"\n",
        "    \"Sono Andrea\"\n",
        "\n",
        "define user greeting\n",
        "    \"Hey there!\"\n",
        "    \"How are you?\"\n",
        "    \"What's up?\"\n",
        "\n",
        "define bot name greeting\n",
        "    \"Hey $name!\"\n",
        "\n",
        "define flow give name\n",
        "    user give name\n",
        "    $name = ...\n",
        "    bot name greeting\n",
        "\n",
        "define flow\n",
        "    user greeting\n",
        "    if not $name\n",
        "        bot ask name\n",
        "    else\n",
        "        bot name greeting\n",
        "\"\"\""
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "LxRmdI6_AjTu"
      },
      "source": [
        "Remove `$name` parameter from the `context`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "_oqkWp9lAmXw"
      },
      "outputs": [],
      "source": [
        "messages = [\n",
        "    {\"role\": \"context\", \"content\": \"\"},\n",
        "    {\"role\": \"assistant\",\n",
        "     \"content\": \"Hi! How are you? Is there anything I can help with?\"},\n",
        "    {\"role\": \"user\", \"content\": \"Hey there!\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "SE_RaWj18TqJ"
      },
      "outputs": [],
      "source": [
        "# initialize rails config\n",
        "config = RailsConfig.from_content(\n",
        "    colang_content=colang_content,\n",
        "    yaml_content=yaml_content\n",
        ")\n",
        "# create rails\n",
        "rails = LLMRails(config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hNZvCksv8WAX",
        "outputId": "ff0f20cb-cef2-4205-b76c-b950008ce7e5"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'role': 'assistant', 'content': \"Hi there! What's your name?\"}"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "res = await rails.generate_async(messages=messages)\n",
        "res"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mYXnob8w8bNs",
        "outputId": "4b65c3bc-a85b-46cd-a989-ee3f8c62062d"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'role': 'assistant', 'content': 'Hey James!'}"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "messages += [\n",
        "    res,\n",
        "    {\"role\": \"user\", \"content\": \"I'm James\"}\n",
        "]\n",
        "res = await rails.generate_async(messages=messages)\n",
        "res"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "rYh_FC-HmK8z"
      },
      "source": [
        "---"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
